import typing as T
from langchain_core.documents import Document
from typing_extensions import override
from io import BytesIO
from jsonfinder import jsonfinder
from pathlib import Path
from . import IFileLoader
from ...utils.AliOssUtils import AliOssUtils
from ....init.ChatClient import get_temp_client
from ...chat import ChatPrompt, MessageType, RoleType, MessageContentFactory
from ...chat.content.InnerData import ImageUrlData


class ImageLoader(IFileLoader):
    def __init__(self): ...

    @override
    def load(
        self,
        blob: T.Union[str, Path, bytes, BytesIO],
        uuid: str,
        file_name: str = None,
        **kwargs,
    ) -> T.List[Document]:
        self.setUuid(uuid)
        blob = self.verify(blob)
        url = AliOssUtils.uploadFile(kwargs.get("object_name"), blob).resp.response.url
        messages = [
            {
                "role": RoleType.USER,
                "content": [
                    MessageContentFactory.create(
                        type=MessageType.IMAGE_URL, image_url=ImageUrlData(url=url)
                    ).model_dump(),
                    MessageContentFactory.create(
                        type=MessageType.TEXT, text=ChatPrompt.imageTemplate
                    ).model_dump(),
                ],
            }
        ]
        openai_client = get_temp_client()
        generator = openai_client.chat.completions.create(
            model="qwen-omni-turbo",
            messages=messages,
            modalities=["text"],
            stream=True,
        )
        content = []
        for chunk in generator:
            if chunk.choices and isinstance(chunk.choices[0].delta.content, str):
                content.append(chunk.choices[0].delta.content)
        content = "".join(content)
        text = next(jsonfinder(content, json_only=True))[2].get("text")
        return [
            Document(
                page_content=text,
                metadata={
                    "file_name": file_name,
                    "type": "image",
                    "url": url,
                    "uuid": self.uuid,
                },
            )
        ]
